from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
from pydantic import BaseModel, Field
from ..settings.settings import Settings
from openai import OpenAI

router = APIRouter()

# -----------------------------
# 定义 Pydantic 数据模型
# -----------------------------
class PoemInfo(BaseModel):
    title: str = Field(description="诗歌标题")
    author: str = Field(description="诗人名字")
    content: str = Field(description="诗歌内容")

# 初始化 PydanticOutputParser
parser = PydanticOutputParser(pydantic_object=PoemInfo)

# -----------------------------
# 定义 Prompt 模板
# -----------------------------
prompt = ChatPromptTemplate.from_template(
    "请提供{name}的诗歌信息。\n{format_instructions}"
)
# -----------------------------
# 本地 Ollama 模型
# -----------------------------
llm = ChatOllama(
    model="qwen3:0.6b",
    base_url="http://172.16.21.38:11436"
)

@router.get("/poem_info")
async def poem_info(name: str = Query(..., description="诗的名称")):
    """调用本地 Ollama 模型，返回诗歌 JSON 信息 (Pydantic 格式)"""
    # 渲染 Prompt（注意要传递 format_instructions）
    prompt_messages = prompt.format_messages(
        name=name,
        format_instructions=parser.get_format_instructions()
    )
    # 调用 LLM
    response = llm.invoke(prompt_messages)
    # 解析为 Pydantic 对象
    try:
        poem = parser.parse(response.content)
    except Exception as e:
        return JSONResponse(
            content={"error": "模型返回内容无法解析为 Pydantic 对象", "details": str(e)},
            status_code=500
        )
    return JSONResponse(content={"poem_info": poem.dict()})

# -----------------------------
# 云端 DeepSeek (OpenAI API)
# -----------------------------
settings = Settings()
DeepSeekUrl = settings.DeepSeekUrl
DeepSeekKey = settings.DeepSeekKey

client = OpenAI(api_key=DeepSeekKey, base_url=DeepSeekUrl)

@router.get("/poem_info_deepseek")
async def poem_info_deepseek(name: str = Query(..., description="诗的名称")):
    """调用云端 DeepSeek 模型，返回诗歌 JSON 信息 (Pydantic 格式)"""
    # 渲染 Prompt
    prompt_messages = prompt.format_messages(
        name=name,
        format_instructions=parser.get_format_instructions()
    )
    # 提取用户消息
    user_content = " ".join([m.content for m in prompt_messages if m.type == "human"])
    # 调用 DeepSeek
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": user_content}],
        temperature=0.7
    )
    # 提取文本
    text_output = response.choices[0].message.content
    # 解析为 Pydantic 对象
    try:
        poem = parser.parse(text_output)
    except Exception as e:
        return JSONResponse(
            content={"error": "云端模型返回内容无法解析为 Pydantic 对象", "details": str(e)},
            status_code=500
        )
    return JSONResponse(content={"poem_info_deepseek": poem.dict()})
